# -*- coding: utf-8 -*-
"""
Created on Tue Oct 26 11:52:07 2021

@author: singularity

训练数据预分析

- 读取数据
- 分析数据特征
- 转换为训练网络需要的格式
"""

import numpy as np
import PIL
import pandas

from sklearn.cluster import KMeans

import matplotlib.pyplot as plt

'''
重要参数
'''

DATA_ROOT_PATH = './tiny_vid/'  # 数据集根路径

LABEL_FILE_SUFFIX = '_gt.txt'   # 标签文件后缀
IMG_FILE_SUFFIX = '.JPEG'

CLASSES = ['bird','car','dog','lizard','turtle']    # 类名称

TRAIN_TEST = [150,30]   # 数据集中，训练集与测试集的分布

'''
数据加载
'''
dataClassNum = len(CLASSES)

'''  图像加载  '''
dataSet_Train = [[] for _ in range(dataClassNum)]
dataSet_Test = [[] for _ in range(dataClassNum)]

for classIndex,className in enumerate(CLASSES,start=0):
    dataPathNow = DATA_ROOT_PATH+className + '/'
    
    # 读取训练集
    for trainId in range(1,TRAIN_TEST[0]+1):
        fileName = ('%06d' % trainId) + IMG_FILE_SUFFIX
        dataSet_Train[classIndex].append(np.array( PIL.Image.open(dataPathNow+fileName) ))
        
    # 读取测试集
    for testId in range(TRAIN_TEST[0]+1,TRAIN_TEST[0]+TRAIN_TEST[1]+1):
        fileName = ('%06d' % testId) + IMG_FILE_SUFFIX
        dataSet_Test[classIndex].append(np.array( PIL.Image.open(dataPathNow+fileName) ))


dataSet_Train = np.array(dataSet_Train)
dataSet_Test = np.array(dataSet_Test)


'''  标签加载  '''
targetSet_Train = [[] for _ in range(dataClassNum)]
targetSet_Test = [[] for _ in range(dataClassNum)]

for classIndex,className in enumerate(CLASSES,start=0):
    labelSet = pandas.read_csv(DATA_ROOT_PATH + className + LABEL_FILE_SUFFIX,
                               header=None,
                               nrows=TRAIN_TEST[0]+TRAIN_TEST[1])
    for i in range(TRAIN_TEST[0]):
        targetSet_Train[classIndex].append(np.array(labelSet[0][i].split(' '), dtype = np.int0))
    for i in range(TRAIN_TEST[0],TRAIN_TEST[0]+TRAIN_TEST[1]):
        targetSet_Test[classIndex].append(np.array(labelSet[0][i].split(' '), dtype = np.int0))

targetSet_Test = np.array(targetSet_Test)
targetSet_Train = np.array(targetSet_Train)

'''  数据+标签样例显示  '''
def show(img:np.ndarray,label:np.ndarray):
    plt.imshow(img)
    ax = plt.gca()
    ax.add_patch(plt.Rectangle((label[1], label[2]),    # 起点坐标（x,y）
                               label[3]-label[1], label[4]-label[2],    # 宽度，高度
                               color="blue", fill=False, linewidth=1))

'''  Anchor分析  '''    
def anchor_analysis(K:int = 6):
        
    # 训练集图像数据，对象宽高提取
    Train_Wid = (targetSet_Train[:,:,3]-targetSet_Train[:,:,1]).flatten()
    Train_Hig = (targetSet_Train[:,:,4]-targetSet_Train[:,:,2]).flatten()
    
    # 显示目标形状数据分布
    analysData = np.stack((Train_Wid,Train_Hig),axis=-1)
    
    plt.figure()
    plt.title("source distribute")
    plt.scatter(analysData[:,0],
                analysData[:,1],
                alpha=0.6,s=4)
    
    # K-Means聚合
    estimator = KMeans(n_clusters=K)
    predClus = estimator.fit_predict(analysData)    # 聚合计算
    
    clustCenters = estimator.cluster_centers_   # 获取聚类中心
    print("聚类中心为：")
    print(clustCenters)
    
    # 聚类结果显示
    plt.figure()
    plt.title("K-Means")
    plt.scatter(analysData[:,0],analysData[:,1],    # 聚类有色点图
                c=predClus,s=4,alpha=0.9)
    
    for px,py in clustCenters:
        plt.text(px, py, "(%.3f,%.3f)" % (px,py),c='r')
        
        plt.scatter(px, py, s=18,c = 'r')
    
    return clustCenters



    